﻿// #include <torch/extension.h> // 包含了 ATen、pybind11等头文件
#include "myop3.h"

#include <ATen/Operators.h>
#include <torch/all.h>
#include <torch/library.h>

extern "C"
{
    /* Creates a dummy empty _C module that can be imported from Python.
 The import from Python will load the .so consisting of this file
 in this extension, so that the TORCH_LIBRARY static initializers
 below are run. */
    PyObject *PyInit_myop3(void)
    {
        static struct PyModuleDef module_def = {
            PyModuleDef_HEAD_INIT,
            "myop3", /* name of module */
            NULL,    /* module documentation, may be NULL */
            -1,      /* size of per-interpreter state of the module,
               or -1 if the module keeps state in global variables. */
            NULL,    /* methods */
        };
        return PyModule_Create(&module_def);
    }
}

namespace myop3
{

    // 前向传播函数
    at::Tensor my_reduce(const at::Tensor &y)
    {
        // y 可以是任意 shape，这里举个最简单的：sum 规约
        return y.sum();  // 你可以替换成自己手写的 SIMD / OpenMP / CUDA 实现
    }

    // 反向传播函数：把上游梯度广播成 y 的形状
    at::Tensor my_reduce_bwd(const at::Tensor &grad_out, const at::Tensor &y) { return grad_out.expand_as(y); }

    // 使用 TorchScript 自定义操作的方式注册正反两个函数，并关联自动微分
    // 这是新版本推荐的方式，替代旧的 PYBIND11_MODULE 方式:cite[6]:cite[7]
    TORCH_LIBRARY(myop3, m)
    {
        // 注册前向函数 "my_reduce"
        m.def("my_reduce(Tensor a)->Tensor");
        // 注册反向函数 "my_reduce_bwd"
        m.def("my_reduce_bwd(Tensor a, Tensor b)->Tensor");

        // 如果你的操作需要支持自动微分，通常不需要像上面那样手动分开注册正反向
        // 更常见的做法是定义一个继承自 torch::autograd::Function 的类
        // 但 TORCH_LIBRARY 方式更底层，适用于与 PyTorch 核心深度集成
    }

    // Registers CPU implementations for mymuladd, mymul, myadd_out
    TORCH_LIBRARY_IMPL(myop3, CPU, m)
    {
        m.impl("my_reduce", &my_reduce);
        m.impl("my_reduce_bwd", &my_reduce_bwd);
    }

}  // namespace myop3

// 注意：在新版本的 PyTorch C++ 扩展实践中，如果使用了 TORCH_LIBRARY，
// 通常就不再需要（或者说不应该同时使用）PYBIND11_MODULE 来暴露同一个函数给
// Python 了。 你应该选择一种方式。对于希望深度集成、高效且能被 TorchScript
// 识别的操作， 推荐使用 TORCH_LIBRARY。

// 如果你仍然需要传统的 Python 绑定（例如为了兼容性，或者需要被直接调用），
// 可以保留 PYBIND11_MODULE，但要注意避免重复定义和链接冲突。
// 但更现代的做法是让 Python 通过 torch.ops.your_namespace.your_operator 来调用
// TORCH_LIBRARY 注册的操作:cite[6]。
#if 0
PYBIND11_MODULE(myop3, m)
{
    // 如果你的构建系统（如 setup.py 或 CMakeLists.txt）定义了 TORCH_EXTENSION_NAME，
    // 例如为 'my_reduce'，那么这个模块名就是 'my_reduce'
    m.def("forward", &my_reduce_fwd, "custom reduce forward");
    m.def("backward", &my_reduce_bwd, "custom reduce backward");
    // 注意：通过 pybind11 这样直接暴露的 'forward' 和 'backward' 函数，
    // 不会自动与 PyTorch 的 autograd 系统关联。它们只是普通的函数。
    // 通常，更常见的做法是定义一个继承 torch.autograd.Function 的 Python 类，
    // 在其 forward 和 backward 方法中分别调用 C++ 的 forward 和 backward 函数:cite[1]:cite[7]。
}
#endif